import time
import math

class MPU6050:
    """
    MPU6050 6轴惯性测量单元驱动类，适用于树莓派
    
    主要功能：
    - 初始化I2C连接
    - 配置传感器
    - 校准传感器
    - 读取原始数据
    - 读取校准后的数据
    - 计算姿态角度
    """
    
    # MPU6050寄存器定义
    PWR_MGMT_1 = 0x6B      # 电源管理寄存器1
    PWR_MGMT_2 = 0x6C      # 电源管理寄存器2
    SMPLRT_DIV = 0x19      # 采样率分频器
    CONFIG = 0x1A          # 配置寄存器
    GYRO_CONFIG = 0x1B     # 陀螺仪配置寄存器
    ACCEL_CONFIG = 0x1C    # 加速度计配置寄存器
    
    # 数据寄存器起始地址
    ACCEL_XOUT_H = 0x3B    # 加速度计X轴高位
    ACCEL_XOUT_L = 0x3C    # 加速度计X轴低位
    ACCEL_YOUT_H = 0x3D    # 加速度计Y轴高位
    ACCEL_YOUT_L = 0x3E    # 加速度计Y轴低位
    ACCEL_ZOUT_H = 0x3F    # 加速度计Z轴高位
    ACCEL_ZOUT_L = 0x40    # 加速度计Z轴低位
    TEMP_OUT_H = 0x41      # 温度传感器高位
    TEMP_OUT_L = 0x42      # 温度传感器低位
    GYRO_XOUT_H = 0x43     # 陀螺仪X轴高位
    GYRO_XOUT_L = 0x44     # 陀螺仪X轴低位
    GYRO_YOUT_H = 0x45     # 陀螺仪Y轴高位
    GYRO_YOUT_L = 0x46     # 陀螺仪Y轴低位
    GYRO_ZOUT_H = 0x47     # 陀螺仪Z轴高位
    GYRO_ZOUT_L = 0x48     # 陀螺仪Z轴低位
    
    # 陀螺仪满量程范围（对应寄存器值）
    GYRO_FS_250 = 0x00     # ±250°/s
    GYRO_FS_500 = 0x08     # ±500°/s
    GYRO_FS_1000 = 0x10    # ±1000°/s
    GYRO_FS_2000 = 0x18    # ±2000°/s
    
    # 加速度计满量程范围（对应寄存器值）
    ACCEL_FS_2G = 0x00     # ±2g
    ACCEL_FS_4G = 0x08     # ±4g
    ACCEL_FS_8G = 0x10     # ±8g
    ACCEL_FS_16G = 0x18    # ±16g
    
    def __init__(self, address=0x68, bus_num=1):
        from smbus2 import SMBus
        """
        初始化MPU6050
        
        Args:
            address: MPU6050的I2C地址，默认0x68
            bus_num: I2C总线编号，树莓派默认1
        """
        self.address = address
        self.bus = SMBus(bus_num)
        
        # 传感器校准值
        self.accel_offset = [0, 0, 0]
        self.gyro_offset = [0, 0, 0]
        
        # 满量程范围及灵敏度
        self.gyro_scale = 131.0  # 默认±250°/s，灵敏度131 LSB/(°/s)
        self.accel_scale = 16384.0  # 默认±2g，灵敏度16384 LSB/g
        
        # 姿态角度
        self.roll = 0.0
        self.pitch = 0.0
        self.yaw = 0.0
        
        # 初始化传感器
        self._initialize()
    
    def _initialize(self):
        """
        初始化MPU6050传感器
        """
        # 唤醒传感器（清除睡眠位）
        self.bus.write_byte_data(self.address, self.PWR_MGMT_1, 0x00)
        time.sleep(0.01)  # 等待传感器稳定
        
        # 设置采样率（1kHz）
        self.bus.write_byte_data(self.address, self.SMPLRT_DIV, 0x07)
        
        # 设置低通滤波器（带宽44Hz）
        self.bus.write_byte_data(self.address, self.CONFIG, 0x03)
        
        # 设置陀螺仪满量程为±250°/s
        self.bus.write_byte_data(self.address, self.GYRO_CONFIG, self.GYRO_FS_250)
        
        # 设置加速度计满量程为±2g
        self.bus.write_byte_data(self.address, self.ACCEL_CONFIG, self.ACCEL_FS_2G)
    
    def set_gyro_range(self, range_setting):
        """
        设置陀螺仪满量程范围
        
        Args:
            range_setting: 可选值：
                          - GYRO_FS_250 (±250°/s)
                          - GYRO_FS_500 (±500°/s)
                          - GYRO_FS_1000 (±1000°/s)
                          - GYRO_FS_2000 (±2000°/s)
        """
        self.bus.write_byte_data(self.address, self.GYRO_CONFIG, range_setting)
        
        # 更新灵敏度
        if range_setting == self.GYRO_FS_250:
            self.gyro_scale = 131.0
        elif range_setting == self.GYRO_FS_500:
            self.gyro_scale = 65.5
        elif range_setting == self.GYRO_FS_1000:
            self.gyro_scale = 32.8
        elif range_setting == self.GYRO_FS_2000:
            self.gyro_scale = 16.4
    
    def set_accel_range(self, range_setting):
        """
        设置加速度计满量程范围
        
        Args:
            range_setting: 可选值：
                          - ACCEL_FS_2G (±2g)
                          - ACCEL_FS_4G (±4g)
                          - ACCEL_FS_8G (±8g)
                          - ACCEL_FS_16G (±16g)
        """
        self.bus.write_byte_data(self.address, self.ACCEL_CONFIG, range_setting)
        
        # 更新灵敏度
        if range_setting == self.ACCEL_FS_2G:
            self.accel_scale = 16384.0
        elif range_setting == self.ACCEL_FS_4G:
            self.accel_scale = 8192.0
        elif range_setting == self.ACCEL_FS_8G:
            self.accel_scale = 4096.0
        elif range_setting == self.ACCEL_FS_16G:
            self.accel_scale = 2048.0
    
    def _read_word_2c(self, reg):
        """
        读取16位有符号数据
        
        Args:
            reg: 寄存器地址
            
        Returns:
            int: 16位有符号数据
        """
        high = self.bus.read_byte_data(self.address, reg)
        low = self.bus.read_byte_data(self.address, reg + 1)
        value = (high << 8) + low
        
        # 转换为有符号整数
        if value >= 0x8000:
            return -((65535 - value) + 1)
        else:
            return value
    
    def get_raw_data(self):
        """
        读取原始传感器数据
        
        Returns:
            dict: 包含原始数据的字典，格式：
                 {
                     'accel_x': 加速度计X轴原始值,
                     'accel_y': 加速度计Y轴原始值,
                     'accel_z': 加速度计Z轴原始值,
                     'gyro_x': 陀螺仪X轴原始值,
                     'gyro_y': 陀螺仪Y轴原始值,
                     'gyro_z': 陀螺仪Z轴原始值,
                     'temp': 温度原始值
                 }
        """
        accel_x = self._read_word_2c(self.ACCEL_XOUT_H)
        accel_y = self._read_word_2c(self.ACCEL_YOUT_H)
        accel_z = self._read_word_2c(self.ACCEL_ZOUT_H)
        gyro_x = self._read_word_2c(self.GYRO_XOUT_H)
        gyro_y = self._read_word_2c(self.GYRO_YOUT_H)
        gyro_z = self._read_word_2c(self.GYRO_ZOUT_H)
        temp = self._read_word_2c(self.TEMP_OUT_H)
        
        return {
            'accel_x': accel_x,
            'accel_y': accel_y,
            'accel_z': accel_z,
            'gyro_x': gyro_x,
            'gyro_y': gyro_y,
            'gyro_z': gyro_z,
            'temp': temp
        }
    
    def get_scaled_data(self):
        """
        读取校准并缩放后的传感器数据
        
        Returns:
            dict: 包含校准后数据的字典，格式：
                 {
                     'accel_x': 加速度计X轴值 (g),
                     'accel_y': 加速度计Y轴值 (g),
                     'accel_z': 加速度计Z轴值 (g),
                     'gyro_x': 陀螺仪X轴值 (°/s),
                     'gyro_y': 陀螺仪Y轴值 (°/s),
                     'gyro_z': 陀螺仪Z轴值 (°/s),
                     'temp': 温度值 (°C)
                 }
        """
        raw_data = self.get_raw_data()
        
        # 校准并缩放数据
        accel_x = (raw_data['accel_x'] - self.accel_offset[0]) / self.accel_scale
        accel_y = (raw_data['accel_y'] - self.accel_offset[1]) / self.accel_scale
        accel_z = (raw_data['accel_z'] - self.accel_offset[2]) / self.accel_scale
        
        gyro_x = (raw_data['gyro_x'] - self.gyro_offset[0]) / self.gyro_scale
        gyro_y = (raw_data['gyro_y'] - self.gyro_offset[1]) / self.gyro_scale
        gyro_z = (raw_data['gyro_z'] - self.gyro_offset[2]) / self.gyro_scale
        
        # 转换温度（公式：°C = (TEMP_OUT - 32768) / 340 + 36.53）
        temp = (raw_data['temp'] - 32768) / 340 + 36.53
        
        return {
            'accel_x': accel_x,
            'accel_y': accel_y,
            'accel_z': accel_z,
            'gyro_x': gyro_x,
            'gyro_y': gyro_y,
            'gyro_z': gyro_z,
            'temp': temp
        }
    
    def calibrate(self, samples=1000):
        """
        校准传感器
        
        Args:
            samples: 校准采样次数，默认1000次
        """
        print(f"正在校准MPU6050，共{samples}次采样...")
        
        # 初始化校准值
        accel_sum = [0, 0, 0]
        gyro_sum = [0, 0, 0]
        
        # 读取多次采样数据
        for i in range(samples):
            raw_data = self.get_raw_data()
            accel_sum[0] += raw_data['accel_x']
            accel_sum[1] += raw_data['accel_y']
            accel_sum[2] += raw_data['accel_z']
            gyro_sum[0] += raw_data['gyro_x']
            gyro_sum[1] += raw_data['gyro_y']
            gyro_sum[2] += raw_data['gyro_z']
            time.sleep(0.001)  # 1ms延迟
        
        # 计算平均值
        self.accel_offset[0] = accel_sum[0] / samples
        self.accel_offset[1] = accel_sum[1] / samples
        self.accel_offset[2] = accel_sum[2] / samples - self.accel_scale  # Z轴应等于1g
        
        self.gyro_offset[0] = gyro_sum[0] / samples
        self.gyro_offset[1] = gyro_sum[1] / samples
        self.gyro_offset[2] = gyro_sum[2] / samples
        
        print(f"校准完成！")
        print(f"加速度计偏移：{self.accel_offset}")
        print(f"陀螺仪偏移：{self.gyro_offset}")
    
    def calculate_angles(self, scaled_data, dt):
        """
        使用互补滤波计算姿态角度
        
        Args:
            scaled_data: 校准后的传感器数据
            dt: 采样时间间隔（秒）
        
        Returns:
            dict: 包含姿态角度的字典，格式：
                 {
                     'roll': 横滚角 (°),
                     'pitch': 俯仰角 (°),
                     'yaw': 偏航角 (°)
                 }
        """
        # 从加速度计计算角度
        accel_roll = math.atan2(scaled_data['accel_y'], scaled_data['accel_z']) * 57.3
        accel_pitch = math.atan2(-scaled_data['accel_x'], math.sqrt(scaled_data['accel_y']**2 + scaled_data['accel_z']**2)) * 57.3
        
        # 从陀螺仪计算角度变化
        gyro_roll_delta = scaled_data['gyro_x'] * dt
        gyro_pitch_delta = scaled_data['gyro_y'] * dt
        gyro_yaw_delta = scaled_data['gyro_z'] * dt
        
        # 互补滤波融合数据（加速度计权重0.02，陀螺仪权重0.98）
        self.roll = 0.98 * (self.roll + gyro_roll_delta) + 0.02 * accel_roll
        self.pitch = 0.98 * (self.pitch + gyro_pitch_delta) + 0.02 * accel_pitch
        self.yaw += gyro_yaw_delta  # 偏航角只能由陀螺仪计算
        
        return {
            'roll': self.roll,
            'pitch': self.pitch,
            'yaw': self.yaw
        }
    
    def get_orientation(self):
        """
        获取当前姿态角度
        
        Returns:
            dict: 包含姿态角度的字典
        """
        return {
            'roll': self.roll,
            'pitch': self.pitch,
            'yaw': self.yaw
        }
    
    def reset_orientation(self):
        """
        重置姿态角度
        """
        self.roll = 0.0
        self.pitch = 0.0
        self.yaw = 0.0
    
    def close(self):
        """
        关闭I2C连接
        """
        self.bus.close()

# 示例用法
if __name__ == "__main__":
    try:
        # 初始化MPU6050，使用默认地址0x68和总线1
        mpu = MPU6050()
        
        # 校准传感器
        mpu.calibrate()
        
        print("MPU6050测试")
        print("按Ctrl+C退出")
        
        # 初始化时间
        last_time = time.time()
        
        while True:
            # 计算采样时间间隔
            current_time = time.time()
            dt = current_time - last_time
            last_time = current_time
            
            # 获取校准后的数据
            scaled_data = mpu.get_scaled_data()
            
            # 计算姿态角度
            angles = mpu.calculate_angles(scaled_data, dt)
            
            # 打印数据
            print(f"\r加速度: X={scaled_data['accel_x']:.2f}g, Y={scaled_data['accel_y']:.2f}g, Z={scaled_data['accel_z']:.2f}g | "
                  f"陀螺仪: X={scaled_data['gyro_x']:.2f}°/s, Y={scaled_data['gyro_y']:.2f}°/s, Z={scaled_data['gyro_z']:.2f}°/s | "
                  f"温度: {scaled_data['temp']:.2f}°C | "
                  f"姿态: 横滚={angles['roll']:.2f}°, 俯仰={angles['pitch']:.2f}°, 偏航={angles['yaw']:.2f}°", end="")
            
            time.sleep(0.01)  # 10ms延迟
            
    except KeyboardInterrupt:
        print("\n测试结束")
    finally:
        # 清理资源
        mpu.close()